

// host photon structure
struct HostPhoton
{
	float3  m_Position;
	uchar4  m_Power;
	uchar2  m_Angles;
	ushort  m_Flags;
};

// device photon structure
struct DevicePhotons
{
	float4* m_Position;
	uchar4* m_Power;
};

struct HashPositions
{
	uint* m_HashValue;
	uint* m_PhotonIdx;
};

// this table is indexed with the hash value and refers to the first value in HashList
struct HashTable
{
	int*  m_FirstIdx;
	uint* m_NumPhotons;
};

struct HashParams
{
	float  m_CellSize;
	float3 m_WorldOrigin;
	float3 m_CellScaleFactor;
	uint3  m_GridSize;
};

__constant__ HashParams g_Params;


__device__ int3 GetGridPos(float4 p_Pos)
{
	int3 gridPos;
	gridPos.x = floorf((p_Pos.x - g_Params.m_WorldOrigin.x) * g_Params.m_CellScaleFactor.x);
	gridPos.y = floorf((p_Pos.y - g_Params.m_WorldOrigin.y) * g_Params.m_CellScaleFactor.y);
	gridPos.z = floorf((p_Pos.z - g_Params.m_WorldOrigin.z) * g_Params.m_CellScaleFactor.z);
// 
	if(gridPos.x < 0 || gridPos.y < 0 || gridPos.z < 0)
		gridPos.x = 0;

	return gridPos;
}

// specify spatial hash table size (should be 2^n)
const uint SpatialHashTableSize = 2048;

//#define USE_SPATIAL_HASHING
#if defined USE_SPATIAL_HASHING
__device__ uint CalcGridHash(int3 p_GridPos)
{
	const uint p1 = 73856093; // some large primes
	const uint p2 = 19349663;
	const uint p3 = 83492791;
	uint n = __umul24(p_GridPos.x, p1) ^ __umul24(p_GridPos.y, p2) ^ __umul24(p_GridPos.z, p3);
	return (n & (SpatialHashTableSize - 1));	// % SpatialHashTableSize if a multiple of 2
}
#else
__device__ uint CalcGridHash(int3 p_GridPos)
{
	return __umul24(__umul24(p_GridPos.z, g_Params.m_GridSize.y), g_Params.m_GridSize.x) + __umul24(p_GridPos.y, g_Params.m_GridSize.x) + p_GridPos.x;
}
#endif

// convert array of structs into struct of arrays
__global__ void ConvertInputData(const HostPhoton* p_Input, DevicePhotons p_Output, uint p_Elements)
{
	uint index = __mul24(blockIdx.x, blockDim.x) + threadIdx.x;
	if (index >= p_Elements) return;

	HostPhoton p;
	p.m_Position = p_Input[index].m_Position;
	p.m_Power = p_Input[index].m_Power;

	// fill device photon list
	p_Output.m_Position[index] = make_float4(p.m_Position, 0);
	p_Output.m_Power[index] = p.m_Power;
}


__global__ void CalcPositionHashes(const float4* p_PhotonPos, HashPositions p_Output, uint p_Elements)
{
	uint index = __mul24(blockIdx.x, blockDim.x) + threadIdx.x;
	if (index >= p_Elements) return;

	// calculate hash value
	float4 curPos = p_PhotonPos[index];
	int3 gridPos = GetGridPos(curPos);
	uint hashValue = CalcGridHash(gridPos);

	// write result into output list
	p_Output.m_HashValue[index] = hashValue;
	p_Output.m_PhotonIdx[index] = index;
}


__device__ int BinarySearch(const uint* p_List, uint p_Size, uint p_Key)
{
	int left, right, midpt;
	left = 0;
	right = p_Size - 1;
	while (left <= right)
	{
		midpt = (left + right) >> 1;
		if(p_Key == p_List[midpt])
			return static_cast<int>(midpt);
		else if (p_Key > p_List[midpt])
			left = midpt + 1;
		else
			right = midpt - 1;
	}

	return -1;
}

__global__ void CreateHashTable(const uint* p_Hashes, HashTable p_Table, uint p_Elements, uint p_NumHashes)
{
	uint index = __mul24(blockIdx.x, blockDim.x) + threadIdx.x;
	if (index >= p_NumHashes) return;

	// perform binary search to find elements with our hash in the sorted hash list
	int pos = BinarySearch(p_Hashes, p_Elements, index);
	
	// locate beginning of hash sequence
	int originalPos = pos;
	if(pos > 0)
	{
		while(p_Hashes[pos] == index && pos >= 0)
			--pos;
		++pos;

		// find end of the sequence
		++originalPos;
		while(p_Hashes[originalPos] == index && originalPos < p_NumHashes)
			++originalPos;
	}

	// store result in table
	p_Table.m_FirstIdx[index] = pos;
	p_Table.m_NumPhotons[index] = originalPos - pos;
}

template<bool ReadPhase>
__global__ void ReorderPhotons(DevicePhotons p_Photons, DevicePhotons p_Buffer, uint* p_NewOrder, uint p_Elements)
{
	uint index = __mul24(blockIdx.x, blockDim.x) + threadIdx.x;
	if (index >= p_Elements) return;

	if(ReadPhase == true)
	{
		float4 pos = p_Photons.m_Position[index];
		uchar4 power = p_Photons.m_Power[index];
		p_Buffer.m_Position[index] = pos;
		p_Buffer.m_Power[index] = power;
	}
	else
	{
		uint oldIndex = p_NewOrder[index];
		float4 pos = p_Buffer.m_Position[oldIndex];
		uchar4 power = p_Buffer.m_Power[oldIndex];
		p_Photons.m_Position[index] = pos;
		p_Photons.m_Power[index] = power;
	}
}

__device__ void InsertPhoton(float* p_SqDist, uchar4* p_Powers, float p_NewDist, uchar4 p_NewPower, uint p_FreePos)
{
	int i = 0;
	
	// find correct position
	for(i = 0; i < p_FreePos && p_NewDist > p_SqDist[i]; ++i)
		;

	for(int k = p_FreePos; k > i; --k)
	{
		p_SqDist[k] = p_SqDist[k - 1];
		p_Powers[k] = p_Powers[k - 1];
	}
	p_SqDist[i] = p_NewDist;
	p_Powers[i] = p_NewPower;
}

// insertion sort as described on http://en.wikipedia.org/wiki/Insertion_sort
__device__ void FinalInsertionSort(float* p_Keys, uchar4* p_Values, uint p_Elements, uint p_MaxSortElements)
{
	float maxKey = p_Keys[0];
	for(int l = 1; l < p_MaxSortElements; ++l)
	{
		float curKey = p_Keys[l];
		uchar4 curValue = p_Values[l];
		if(curKey < maxKey)
		{
			int i;
			for(i = 0; i < l && p_Keys[i] <= curKey; ++i)
				;

			for(int j = l; j > i; j--)
			{
				p_Keys[j] = p_Keys[j - 1];
				p_Values[j] = p_Values[j - 1];
			}
			p_Keys[i] = curKey;
			p_Values[i] = curValue;
		}
		maxKey = curKey;
	}

	for(int l = p_MaxSortElements; l < p_Elements; ++l)
	{
		float curKey = p_Keys[l];
		uchar4 curValue = p_Values[l];
		if(curKey < maxKey)
		{
			int i;
			for(i = 0; i < p_MaxSortElements && p_Keys[i] <= curKey; ++i)
				;

			for(int j = p_MaxSortElements - 1; j > i; --j)
			{
				p_Keys[j] = p_Keys[j - 1];
				p_Values[j] = p_Values[j - 1];
			}
			p_Keys[i] = curKey;
			p_Values[i] = curValue;
			maxKey = p_Keys[p_MaxSortElements - 1];
		}
	}
}

__global__ void PhotonSearch(DevicePhotons p_Photons, HashTable p_HT, float4* p_Query, 
							 uint p_NumPoints, float p_MaxSqDist, uint p_MaxPhotons, float* p_SqDist, uchar4* p_Powers)
{
	uint cellIndex = threadIdx.x + __mul24(threadIdx.y, 3) + __mul24(threadIdx.z, 9);
	uint qIndex = __mul24(__mul24(gridDim.y, gridDim.x), blockIdx.z) + __mul24(gridDim.x, blockIdx.y) + blockIdx.x;

	// get shared memory location
	extern __shared__ float sharedMem[];
	float* sqDist = (float*)&sharedMem[p_MaxPhotons * cellIndex];
	uchar4* powers = (uchar4*)&sqDist[p_MaxPhotons * 27];

	if(qIndex < p_NumPoints)
	{
		// init square distances memory
		for(int i = 0; i < p_MaxPhotons; ++i) sqDist[i] = 100000.0f;

		// calc cells position
		float4 curPos = p_Query[qIndex];
		float4 offsetPos = curPos;
		curPos.x += (-g_Params.m_CellSize) + (g_Params.m_CellSize * threadIdx.x);
		curPos.y += (-g_Params.m_CellSize) + (g_Params.m_CellSize * threadIdx.y);
		curPos.z += (-g_Params.m_CellSize) + (g_Params.m_CellSize * threadIdx.z);

		// get cell hash
		int3 gridPos = GetGridPos(offsetPos);
		uint hashValue = CalcGridHash(gridPos);

		// process photons
		float curMaxsqDist = 1000000.0f;
		uint lastResultIdx = p_MaxPhotons - 1;
		uint curNumPhotons = 0;
		int curPhoton = p_HT.m_FirstIdx[hashValue];
		uint numPhotons = p_HT.m_NumPhotons[hashValue];
		for(uint i = 0; i < numPhotons; ++i)
		{
			// calc photon distance
			float4 distVec = p_Photons.m_Position[curPhoton] - curPos;
			float curSqDist = distVec.x * distVec.x + distVec.y * distVec.y + distVec.z * distVec.z;

			if(curSqDist < curMaxsqDist && curSqDist < p_MaxSqDist)
			{
				InsertPhoton(sqDist, powers, curSqDist, p_Photons.m_Power[curPhoton], curNumPhotons);
				if(curNumPhotons < lastResultIdx) ++curNumPhotons;
				curMaxsqDist = sqDist[lastResultIdx];
			}
			++curPhoton;
		}

		float* resultSqDist = p_SqDist + qIndex * p_MaxPhotons * 27 + cellIndex * p_MaxPhotons;
		uchar4* resultPower = p_Powers + qIndex * p_MaxPhotons * 27 + cellIndex * p_MaxPhotons;
		for(int i = 0; i < p_MaxPhotons; ++i)
		{
			resultSqDist[i] = sqDist[i];
			resultPower[i] = powers[i];
		}
	}

// 
// 	__syncthreads();
// 	if(threadIdx.x == 0)
// 	{
// 		FinalInsertionSort(sqDist, (uint*)powers, p_MaxPhotons * 27, p_MaxPhotons);
// 		
// 		// write back to result
// 		float* resultSqDist = p_SqDist + qIndex * (p_MaxPhotons);
// 		uchar4* resultPower = p_Powers + qIndex * (p_MaxPhotons);
// 		for(int i = 0; i < p_MaxPhotons; ++i)
// 		{
// 			resultSqDist[i] = sqDist[i];
// 			resultPower[i] = powers[i];
// 		}
// 	}
}


__global__ void FinalizeResults(float* p_SqDist, uchar4* p_Powers, float* p_ResultSq, uchar4* p_ResultPower, uint p_MaxPhotons, uint p_Elements)
{
	uint index = __mul24(blockIdx.x, blockDim.x) + threadIdx.x;
	if (index >= p_Elements) return;

	float* sqDist = p_SqDist + p_MaxPhotons * 27 * index;
	uchar4* powers = p_Powers + p_MaxPhotons * 27 * index;
	FinalInsertionSort(sqDist, (uchar4*)powers, p_MaxPhotons * 27, p_MaxPhotons);
	
	// write back to result
	float* resultSqDist = p_ResultSq + index * (p_MaxPhotons);
	uchar4* resultPower = p_ResultPower + index * (p_MaxPhotons);
	for(int i = 0; i < p_MaxPhotons; ++i)
	{
		resultSqDist[i] = sqDist[i];
		resultPower[i] = powers[i];
	}
}

__constant__ float4 g_CellOffsets[] = 
{
	{0, 0, 0, 0},

	{1, 0, 0, 0},
	{1, 0, 1, 0},
	{1, 0, -1, 0},

	{1, 1, 0, 0},
	{1, 1, 1, 0},
	{1, 1, -1, 0},

	{1, -1, 0, 0},
	{1, -1, 1, 0},
	{1, -1, -1, 0},

	{-1, 0, 0, 0},
	{-1, 0, 1, 0},
	{-1, 0, -1, 0},

	{-1, 1, 0, 0},
	{-1, 1, 1, 0},
	{-1, 1, -1, 0},

	{-1, -1, 0, 0},
	{-1, -1, 1, 0},
	{-1, -1, -1, 0},

	{0, 0, 1, 0},
	{0, -1, 1, 0},
	{0, 1, 1, 0},

	{0, 0, -1, 0},
	{0, -1, -1, 0},
	{0, 1, -1, 0},

	{0, 1, 0, 0},
	{0, -1, 0, 0}
};

__global__ void PhotonSearch2(DevicePhotons p_Photons, HashTable p_HT, float4* p_Query, 
							 uint p_NumPoints, float p_MaxSqDist, uint p_MaxPhotons, float* p_SqDist, uchar4* p_Powers)
{
	uint qIndex = __mul24(blockIdx.x, blockDim.x) + threadIdx.x;

	// get shared memory locations
	extern __shared__ float sharedMem[];
	float* sqDist = (float*)&sharedMem[2 * p_MaxPhotons * threadIdx.x];
	uchar4* powers = (uchar4*)&sqDist[p_MaxPhotons];

	if(qIndex < p_NumPoints)
	{
		// get current query position
		float4 curPos = p_Query[qIndex];

		// init square distances memory
		for(int i = 0; i < p_MaxPhotons; ++i) sqDist[i] = 100000.0f;

		float curMaxsqDist = 1000000.0f;
		uint lastResultIdx = p_MaxPhotons - 1;
		uint curNumPhotons = 0;
		for(int curCell = 0; curCell < 27; ++curCell)
		{
			// calc cells position
			float4 offsetPos = curPos + g_CellOffsets[curCell];

			// get cell hash
			int3 gridPos = GetGridPos(offsetPos);
			uint hashValue = CalcGridHash(gridPos);

			// process photons
			int curPhoton = p_HT.m_FirstIdx[hashValue];
			uint numPhotons = p_HT.m_NumPhotons[hashValue];
			for(uint i = 0; i < numPhotons; ++i)
			{
				// calc photon distance
				float4 distVec = p_Photons.m_Position[curPhoton] - curPos;
				float curSqDist = distVec.x * distVec.x + distVec.y * distVec.y + distVec.z * distVec.z;

				if(curSqDist < curMaxsqDist && curSqDist < p_MaxSqDist)
				{
					InsertPhoton(sqDist, powers, curSqDist, p_Photons.m_Power[curPhoton], curNumPhotons);
					if(curNumPhotons < lastResultIdx) ++curNumPhotons;
					curMaxsqDist = sqDist[lastResultIdx];
				}
				++curPhoton;
			}
		}

		float* resultSqDist = p_SqDist + qIndex * p_MaxPhotons;
		uchar4* resultPower = p_Powers + qIndex * p_MaxPhotons;
		for(int i = 0; i < p_MaxPhotons; ++i)
		{
			resultSqDist[i] = sqDist[i];
			resultPower[i] = powers[i];
		}
	}
}
